import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import torch
from torch import nn
import torch.nn.functional as F
from tqdm import tqdm
plt.switch_backend('Qt5Agg')

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(1,512)
        self.fc2 = nn.Linear(512,256)
        self.fc3 = nn.Linear(256,1)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x
if __name__ == '__main__':

    # 随机种子
    np.random.seed(42)
    # 准备数据
    x = np.random.rand(100,1)*10
    # 转为张量
    x = torch.tensor(x,dtype=torch.float32)
    y = np.sin(x) + 0.2*np.random.rand(100,1)
    # 转为张量
    target = torch.tensor(y,dtype=torch.float32)
    # 测试数据
    x_test = torch.FloatTensor(np.linspace(0,10,100).reshape(-1,1))

    # 搭建网络结构
    # 初始化
    net = Net()
    # 损失函数
    criterion = nn.MSELoss()
    # 学习率
    lr = 0.001

    # 训练
    for epoch in tqdm(range(70000),total=70000):

        # 前向传播
        out = net(x)
        # 计算损失
        loss = criterion(out,target)
        # 梯度清零
        net.zero_grad()

        # 反向传播
        loss.backward()
        # 更新参数
        for param in net.parameters():
            param.data -= lr*param.grad
        if epoch % 1000 == 0:
            with torch.no_grad():
                # 测试
                out = net(x_test)
            plt.ion()
            plt.cla()
        #     散点图
            plt.plot(x,y,'.')
        #     折线图
            plt.plot(x_test,out)
            plt.pause(0.01)
            plt.ioff()
            plt.savefig('test.png')
